# -*- coding: utf-8 -*-

import sim
import os
import time
import pandas as pd

LAG = int(input('Input LAG:'))


def text2dict(text, s):
    d = {}
    for word in s:
        d[word] = text.count(word)
    return d


def ifcompute(code1, code2, year1, year2, ind1, ind2):
    bool1 = (code1 != code2)
    bool2 = (year2 - year1 == LAG)
    bool3 = (ind1 == ind2)
    return bool1 & bool2 & bool3


def ifcontinue(year1, year2, ind1, ind2):
    bool1 = (abs(year2 - year1) <= abs(LAG))
    bool2 = (ind1 == ind2)
    return bool1 & bool2


def compute(df, i, forward=True):
    ret1_each = []
    ret2_each = []
    j = i
    ifnext = True
    while ifnext:
        if forward == True:
            j = j + 1
        elif forward == False:
            j = j - 1
        else:
            raise ValueError
        try:
            code1 = (6 - len(str(int(df.stkcd[i])))) * '0' + str(
                int(df.stkcd[i]))
            code2 = (6 - len(str(int(df.stkcd[j])))) * '0' + str(
                int(df.stkcd[j]))
            year1 = df.year[i]
            year2 = df.year[j]
            ind1 = df.ind[i]
            ind2 = df.ind[j]
        except KeyError:
            break
        if ifcontinue(year1, year2, ind1, ind2):
            if ifcompute(code1, code2, year1, year2, ind1, ind2):
                file1 = '{}-{}.txt'.format(year1, code1)
                file2 = '{}-{}.txt'.format(year2, code2)
                text1 = sim.readtext(path, file1)
                text2 = sim.readtext(path, file2)
                wordset1 = set(df.WordSet_Full[i].replace("'", '').replace(
                    ',', '')[1:][:-1].split())
                wordset2 = set(df.WordSet_Full[j].replace("'", '').replace(
                    ',', '')[1:][:-1].split())
                ws1 = set(df.WordSet[i].replace("'", '').replace(
                    ',', '')[1:][:-1].split())
                ws2 = set(df.WordSet[j].replace("'", '').replace(
                    ',', '')[1:][:-1].split())
                d1 = text2dict(text1, ws1)
                d2 = text2dict(text2, ws2)
                v1, v2 = sim.dict2vec(d1, d2)
                sim1 = sim.jaccard(wordset1, wordset2)
                if sim1 != None:
                    ret1_each.append(sim1)
                sim2 = sim.cos_sim(v1, v2)
                if sim2 != None:
                    ret2_each.append(sim2)
        else:
            ifnext = False
    return ret1_each, ret2_each


if __name__ == "__main__":

    path = input("请输入待检索文件目录：")
    file_list = [
        x for x in os.listdir(path) if os.path.splitext(x)[-1] == '.txt'
    ]
    print("{:}个待检索文件读取成功!".format(len(file_list)))

    csvpath = input("input csv file path:")
    df = pd.read_csv(csvpath)
    df = df.sort_values(by=["ind", "year"], ascending=(LAG >= 0))
    df.index = range(len(df))
    ret1 = []
    ret2 = []

    ret_path = input("input output file path:")

    start_time = time.time()
    for i in range(len(df)):
        ret1_each_f, ret2_each_f = compute(df, i, forward=True)
        ret1_each_b, ret2_each_b = compute(df, i, forward=False)
        ret1_each = ret1_each_f + ret1_each_b
        ret2_each = ret2_each_f + ret2_each_b

        if len(ret1_each) != 0:
            ret1.append(sum(ret1_each) / len(ret1_each))
        else:
            ret1.append(None)
        if len(ret2_each) != 0:
            ret2.append(sum(ret2_each) / len(ret2_each))
        else:
            ret2.append(None)
        end_time = time.time()
        remain_time = (end_time - start_time) / (i + 1) * (len(df) - (i + 1))
        print(
            "\r{:>6.2f}% Done, Time Remained: {:04}:{:02}:{:02}".format(
                (i + 1) * 100 / len(df), int(remain_time // 3600),
                int(remain_time // 60 % 60), int(remain_time % 60)),
            end="")

    df = df.drop(['WordSet'], axis=1)
    df = df.drop(['WordSet_Full'], axis=1)
    df['Sim1'] = ret1
    df['Sim2'] = ret2

    df.to_csv(ret_path, index=False)
